# _*_ coding:utf-8 _*_
"""
@author:lx
@file: segement2
@contact: lixiang-929@outlook.com
@time: 2019/03/19
"""
# coding=utf-8
import cv2
import numpy as np
import os
import time

data_path = ".//data//"
result_path = ".//result//"


def color_segment(data_path, result_path):
    data_list = os.listdir(data_path)
    for data in data_list:
        start_time = time.time()
        image_name = os.path.join(data_path, data)
        image = cv2.imread(image_name)
        image_width = image.shape[1]
        image_height = image.shape[0]

        r2_g_b_image = np.zeros([image_height, image_width])
        g_b_image = np.zeros([image_height, image_width])

        image1 = np.zeros([image_height, image_width])  # gray_image which is heat map of red
        image2 = np.zeros([image_height, image_width])  # gray_image which is heat map of yellow

        b_channel, g_channel, r_channel = cv2.split(image)

        for j in range(image_height):
            for i in range(image_width):
                if (r_channel[j, i] *2) > (g_channel[j, i] + b_channel[j,i]):
                    image1[j,i] = 2 * r_channel[j, i] - g_channel[j, i] - b_channel[j,i]
                else:
                    image1[j,i] = 0
                image2[j,i] = (g_channel[j,i] + b_channel[j,i])/2

        # for j in range(image_height):
        #     for i in range(image_width):
        #         if rg_image[j, i] > gb_image[j, i]:
        #             red_image[j, i] = rg_image[j, i] - gb_image[j, i]
        #         else:
        #             yellow_image[j, i] = gb_image[j, i] - rg_image[j, i]

        # cv2.imshow("red_image", red_image)
        file_name, file_format = os.path.splitext(data)
        image1_name = os.path.join(result_path, file_name + "_image1.jpg")
        image2_name = os.path.join(result_path, file_name + "_image2.jpg")
        cv2.imwrite(image1_name, image1)
        cv2.imwrite(image2_name, image2)
        print("{0:.3f}, {1:s}".format((time.time() - start_time), file_name))
        cv2.waitKey(10)


if __name__ == '__main__':
    color_segment(data_path, result_path)